package za.ac.sun.cs.distrib;

import gov.nasa.jpf.Config;
import gov.nasa.jpf.PropertyListenerAdapter;
import gov.nasa.jpf.jvm.bytecode.IfInstruction;
import gov.nasa.jpf.jvm.bytecode.JVMInvokeInstruction;
import gov.nasa.jpf.search.Search;
import gov.nasa.jpf.symbc.SymbolicInstructionFactory;
import gov.nasa.jpf.symbc.bytecode.BytecodeUtils;
import gov.nasa.jpf.symbc.bytecode.IFEQ;
import gov.nasa.jpf.symbc.bytecode.IFGE;
import gov.nasa.jpf.symbc.bytecode.IFGT;
import gov.nasa.jpf.symbc.bytecode.IFLE;
import gov.nasa.jpf.symbc.bytecode.IFLT;
import gov.nasa.jpf.symbc.bytecode.IFNE;
import gov.nasa.jpf.symbc.bytecode.IF_ICMPEQ;
import gov.nasa.jpf.symbc.bytecode.IF_ICMPGE;
import gov.nasa.jpf.symbc.bytecode.IF_ICMPGT;
import gov.nasa.jpf.symbc.bytecode.IF_ICMPLE;
import gov.nasa.jpf.symbc.bytecode.IF_ICMPLT;
import gov.nasa.jpf.symbc.bytecode.IF_ICMPNE;
import gov.nasa.jpf.symbc.numeric.Comparator;
import gov.nasa.jpf.symbc.numeric.IntegerExpression;
import gov.nasa.jpf.symbc.numeric.PCChoiceGenerator;
import gov.nasa.jpf.symbc.numeric.PathCondition;
import gov.nasa.jpf.symbc.numeric.SymbolicInteger;
import gov.nasa.jpf.vm.ChoiceGenerator;
import gov.nasa.jpf.vm.ClassInfo;
import gov.nasa.jpf.vm.Instruction;
import gov.nasa.jpf.vm.MethodInfo;
import gov.nasa.jpf.vm.StackFrame;
import gov.nasa.jpf.vm.ThreadInfo;
import gov.nasa.jpf.vm.VM;

import java.lang.management.ManagementFactory;
import java.lang.management.ThreadMXBean;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;

import org.apfloat.Apint;

import redis.clients.jedis.Jedis;
import za.ac.sun.cs.distrib.logging.LogWriter;
import za.ac.sun.cs.green.expr.IntVariable;
import za.ac.sun.cs.green.util.Reporter;

public class TreeCountListener extends PropertyListenerAdapter {

	private static final int RANDOM_CHOICE = 0;

	private static final int FATTEST_CHOICE = 1;
	
	private static final int LEANEST_CHOICE = 2;
	
	private static final int WEIGHED_CHOICE = 3;
	
	private final int strategy;

	private final Jedis db;

	private Search search = null;

	private final static Random random = new Random();

	private final ThreadMXBean bean = ManagementFactory.getThreadMXBean();

	private Apint totalPaths = Apint.ZERO;

	private Apint totalOverlaps = Apint.ZERO;
	
	private Apint totalValues = Apint.ZERO;
	
	private Apint totalUpdates = Apint.ZERO;
	
	private Apint totalChoices = Apint.ZERO;
	
	private Apint totalDeadends = Apint.ZERO;
	
	public TreeCountListener(Config config) {
		String redisHost = config.getString("search.TreeListener.redis.host", "localhost");
		int redisPort = config.getInt("search.TreeListener.redis.port", 7777);
		db = new Jedis(redisHost, redisPort);
		String id = config.getString("search.name");
		if (id == null) {
			id = config.getString("target");
		}
		LogWriter.initialize(config, id);
		String s = config.getString("search.TreeListener.strategy", "random");
		if (s.equals("fattest")) {
			strategy = FATTEST_CHOICE;
		} else if (s.equals("leanest")) {
			strategy = LEANEST_CHOICE;
		} else if (s.equals("weighted")) {
			strategy = WEIGHED_CHOICE;
		} else {
			strategy = RANDOM_CHOICE;
		}
	}

	@Override
	public void searchStarted(Search search) {
		this.search = search;
	}

	@Override
	public void choiceGeneratorAdvanced(VM vm, ChoiceGenerator<?> currentCG) {
		if (currentCG instanceof PCChoiceGenerator) {
			if (currentCG.getTotalNumberOfChoices() > 1) {
				Instruction instruction = vm.getInstruction();
				ThreadInfo threadInfo = vm.getCurrentThread();
				if (instruction instanceof IfInstruction) {
					try {
						if (computeBranchPC((IfInstruction) instruction,
								threadInfo, currentCG)) {
							((PCChoiceGenerator) currentCG).select(0);
						} else {
							((PCChoiceGenerator) currentCG).select(1);
						}
					} catch (TreeException x) {
						x.printStackTrace();
						vm.ignoreState();
					}
				}
			}
		}
	}

	private static final Map<Class<? extends IfInstruction>, Comparator> mapInstruction = new HashMap<>();

	static {
		mapInstruction.put(IFEQ.class, Comparator.EQ);
		mapInstruction.put(IFNE.class, Comparator.NE);
		mapInstruction.put(IFLT.class, Comparator.LT);
		mapInstruction.put(IFLE.class, Comparator.LE);
		mapInstruction.put(IFGT.class, Comparator.GT);
		mapInstruction.put(IFGE.class, Comparator.GE);
	}

	private static final Map<Class<? extends IfInstruction>, Comparator> complInstruction = new HashMap<>();

	static {
		complInstruction.put(IFEQ.class, Comparator.NE);
		complInstruction.put(IFNE.class, Comparator.EQ);
		complInstruction.put(IFLT.class, Comparator.GE);
		complInstruction.put(IFLE.class, Comparator.GT);
		complInstruction.put(IFGT.class, Comparator.LE);
		complInstruction.put(IFGE.class, Comparator.LT);
	}

	private static final Map<Class<? extends IfInstruction>, Comparator> dMapInstruction = new HashMap<>();
	
	static {
		dMapInstruction.put(IF_ICMPEQ.class, Comparator.EQ);
		dMapInstruction.put(IF_ICMPNE.class, Comparator.NE);
		dMapInstruction.put(IF_ICMPLT.class, Comparator.LT);
		dMapInstruction.put(IF_ICMPLE.class, Comparator.LE);
		dMapInstruction.put(IF_ICMPGT.class, Comparator.GT);
		dMapInstruction.put(IF_ICMPNE.class, Comparator.NE);
	}
	
	private static final Map<Class<? extends IfInstruction>, Comparator> dComplInstruction = new HashMap<>();
	
	static {
		dComplInstruction.put(IF_ICMPEQ.class, Comparator.NE);
		dComplInstruction.put(IF_ICMPNE.class, Comparator.EQ);
		dComplInstruction.put(IF_ICMPLT.class, Comparator.GE);
		dComplInstruction.put(IF_ICMPLE.class, Comparator.GT);
		dComplInstruction.put(IF_ICMPGT.class, Comparator.LE);
		dComplInstruction.put(IF_ICMPGE.class, Comparator.LT);
	}
	
	private boolean computeBranchPC(IfInstruction instruction,
			ThreadInfo threadInfo, ChoiceGenerator<?> currentCG)
			throws TreeException {
		ChoiceGenerator<?> previousCG = currentCG.getPreviousChoiceGenerator();
		while (!((previousCG == null) || (previousCG instanceof PCChoiceGenerator))) {
			previousCG = previousCG.getPreviousChoiceGenerator();
		}
		PathCondition previousPC = (previousCG == null) ? new PathCondition()
				: ((PCChoiceGenerator) previousCG).getCurrentPC();
		PathCondition ifPC = previousPC.make_copy();
		PathCondition elsePC = previousPC.make_copy();
		StackFrame stackFrame = threadInfo.getTopFrame();
		
		if (mapInstruction.containsKey(instruction.getClass())) {
			IntegerExpression condition = (IntegerExpression) stackFrame.getOperandAttr();
			ifPC._addDet(mapInstruction.get(instruction.getClass()), condition, 0);
			elsePC._addDet(complInstruction.get(instruction.getClass()), condition, 0);
		} else if (dMapInstruction.containsKey(instruction.getClass())) {
			int v2 = stackFrame.peek();
			int v1 = stackFrame.peek(1);
			IntegerExpression sym_v1 = (IntegerExpression) stackFrame.getOperandAttr(1);
			IntegerExpression sym_v2 = (IntegerExpression) stackFrame.getOperandAttr(0);
			if (sym_v1 != null) {
				if (sym_v2 != null) { // both are symbolic values
					ifPC._addDet(dMapInstruction.get(instruction.getClass()), sym_v1, sym_v2);
					elsePC._addDet(dComplInstruction.get(instruction.getClass()), sym_v1, sym_v2);
				} else {
					ifPC._addDet(dMapInstruction.get(instruction.getClass()), sym_v1, v2);
					elsePC._addDet(dComplInstruction.get(instruction.getClass()), sym_v1, v2);
				}
			} else {
				ifPC._addDet(dMapInstruction.get(instruction.getClass()), v1, sym_v2);
				elsePC._addDet(dComplInstruction.get(instruction.getClass()), v1, sym_v2);
			}
		}

		Apint ifCount = getCount(ifPC);
		Apint elseCount = getCount(elsePC);

		String ifData = cleanse(ifPC);
		String elseData = cleanse(elsePC);
		int ifComparison = ifCount.compareTo(Apint.ZERO);
		int elseComparison = elseCount.compareTo(Apint.ZERO);
		if ((ifComparison <= 0) && (elseComparison <= 0)) {
			LogWriter.nochoice(ifData, elseData);
			totalDeadends = totalDeadends.add(Apint.ONE);
			if (ifData.contains("&&")) {
				((TreeSearch) search).resetSearch();
			} else {
				search.terminate();
			}
			return false;
		} else if (ifComparison <= 0) {
			totalChoices = totalChoices.add(Apint.ONE);
			LogWriter.choice(elseData, elseCount, ifData, ifCount);
			return true;
		} else if (elseComparison <= 0) {
			totalChoices = totalChoices.add(Apint.ONE);
			LogWriter.choice(ifData, ifCount, elseData, elseCount);
			return false;
		} else if (strategy == FATTEST_CHOICE) {
			totalChoices = totalChoices.add(Apint.ONE);
			int r = ifCount.compareTo(elseCount);
			if ((r == 0) ? random.nextBoolean() : (r < 0)) {
				LogWriter.choice(elseData, elseCount, ifData, ifCount);
				return true;
			} else {
				LogWriter.choice(ifData, ifCount, elseData, elseCount);
				return false;
			}
		} else if (strategy == LEANEST_CHOICE) {
			totalChoices = totalChoices.add(Apint.ONE);
			int r = ifCount.compareTo(elseCount);
			if ((r == 0) ? random.nextBoolean() : (r > 0)) {
				LogWriter.choice(elseData, elseCount, ifData, ifCount);
				return true;
			} else {
				LogWriter.choice(ifData, ifCount, elseData, elseCount);
				return false;
			}
		} else if (strategy == WEIGHED_CHOICE) {
			totalChoices = totalChoices.add(Apint.ONE);
			double weight = ifCount.divide(elseCount.add(ifCount)).doubleValue();
			if (random.nextDouble() > weight) {
				LogWriter.choice(elseData, elseCount, ifData, ifCount);
				return true;
			} else {
				LogWriter.choice(ifData, ifCount, elseData, elseCount);
				return false;
			}
		} else {
			totalChoices = totalChoices.add(Apint.ONE);
			if (random.nextBoolean()) {
				LogWriter.choice(elseData, elseCount, ifData, ifCount);
				return true;
			} else {
				LogWriter.choice(ifData, ifCount, elseData, elseCount);
				return false;
			}
		}
	}

	@Override
	public void stateAdvanced(Search search) {
		if (search.isEndState()) {
			searchUpdate(search);
		}
	}

	@Override
	public void searchConstraintHit(Search search) {
		searchUpdate(search);
	}

	private void searchUpdate(Search search) {
		VM vm = search.getVM();
		PCChoiceGenerator previousCG = vm
				.getLastChoiceGeneratorOfType(PCChoiceGenerator.class);
		PathCondition previousPC = (previousCG == null) ? new PathCondition() : previousCG.getCurrentPC().make_copy();
		Apint count = (Apint) previousPC.getInstance().request("count");
		while (previousPC.count() > 0) {
			String k = "KOUNT:" + cleanse(previousPC);
			if (db.exists(k) && updateCount(previousPC, count)) {
				break;
			}
			previousPC.removeHeader();
		}
		if (previousPC.count() == 0) {
			totalValues = totalValues.add(count);
		}
		totalPaths = totalPaths.add(Apint.ONE);
		long spent = bean.isCurrentThreadCpuTimeSupported() ? bean.getCurrentThreadCpuTime() : 0L;
		LogWriter.stats(totalPaths, totalValues, totalOverlaps, totalUpdates, totalChoices, totalDeadends, spent);
		BytecodeUtils.clearSymVarCounter();
	}

	private Apint getCount(PathCondition pathCondition) {
		Apint n = Apint.ZERO;
		String x = cleanse(pathCondition);
		String k = "KOUNT:" + x;
		String s = db.get(k);
		if (s == null) {
			n = (Apint) pathCondition.getInstance().request("count");
			db.setnx(k, n.toString());
			LogWriter.newcount(x, n);
		} else {
			n = new Apint(s);
		}
		return n;
	}

	/**
	 * Decrement the count associated with a given path condition.
	 * 
	 * @param pathCondition
	 *            the path condition whose count is decremented
	 * @param delta
	 *            the amount by which the count is decremented
	 * @return true if and only if the count after the decrement is zero or less
	 */
	private boolean updateCount(PathCondition pathCondition, Apint delta) {
		String x = cleanse(pathCondition);
		String k = "KOUNT:" + x;
		String l = "KLOCK:" + x;
		db.setnx(l, "0");
		String v = "1";
		while (v.equals("1")) { // BUSY WAITING !!!
			v = db.getSet(l, "1");
			// SLEEP ???
		}
		String s = db.get(k);
		assert s != null;
		Apint n = new Apint(s);
		Apint m = n.subtract(delta);
		db.set(k, m.toString());
		db.set(l, "0");
		if (m.compareTo(Apint.ZERO) < 0) {
			LogWriter.conflict(x, n, delta, m);
			totalOverlaps = totalOverlaps.add(Apint.ONE);
			return true;
		} else {
			LogWriter.update(x, n, delta, m);
			totalUpdates = totalUpdates.add(Apint.ONE);
			return false;
		}
	}

	private String cleanse(PathCondition pathCondition) {
		return pathCondition.toString()
				.replaceAll("\n", " ")
				.replaceAll("_[0-9][0-9]*_SYMINT", "")
				.replaceAll("constraint # = [0-9][0-9]* ", "")
				.replaceAll("CONST_", "");
	}

	@Override
	public void instructionExecuted(VM vm, ThreadInfo currentThread,
			Instruction nextInstruction, Instruction executedInstruction) {
		if (!vm.getSystemState().isIgnored() && BounderService.hasNoVariables()) {
			Instruction insn = executedInstruction;
			if (insn instanceof JVMInvokeInstruction) {
				Config conf = vm.getConfig();
				ThreadInfo ti = currentThread;
				StackFrame sf = ti.getTopFrame();
				JVMInvokeInstruction md = (JVMInvokeInstruction) insn;
				MethodInfo mi = md.getInvokedMethod();
				ClassInfo ci = mi.getClassInfo();
				String className = ci.getName();
				String methodName = md.getInvokedMethodName();
				int numberOfArgs = md.getArgumentValues(ti).length;
				if ((BytecodeUtils.isClassSymbolic(conf, className, mi,
						methodName))
						|| BytecodeUtils.isMethodSymbolic(conf,
								mi.getFullName(), numberOfArgs, null)) {
					for (int i = 0; i <= numberOfArgs; i++) {
						Object o = sf.getLocalAttr(i);
						if ((o != null) && (o instanceof SymbolicInteger)) {
							SymbolicInteger j = (SymbolicInteger) o;
							BounderService.addVariable(new IntVariable(j
									.getName(), j, j._min, j._max));
							LogWriter.newvar(j.getName(), j._min, j._max);
						}
					}
				}
			}
		}
	}

	@Override
	public void searchFinished(Search search) {
		SymbolicInstructionFactory.greenSolver.shutdown();
		SymbolicInstructionFactory.greenSolver.report(new Reporter() {
			@Override
			public void report(String context, String message) {
				LogWriter.greenStats(context + ":: " + message);
			}
		});
		long spent = bean.isCurrentThreadCpuTimeSupported() ? bean.getCurrentThreadCpuTime() : 0L;
		LogWriter.done(spent);
	}

}
